package im.composer.audio.outputstream;

import java.io.IOException;

import javax.sound.sampled.AudioFormat;

import org.tritonus.dsp.interfaces.FloatSampleWriter;
import org.tritonus.share.sampled.FloatSampleBuffer;
import org.tritonus.share.sampled.file.TDataOutputStream;
import org.xiph.libogg.ogg_packet;
import org.xiph.libogg.ogg_page;
import org.xiph.libogg.ogg_stream_state;
import org.xiph.libvorbis.vorbis_block;
import org.xiph.libvorbis.vorbis_comment;
import org.xiph.libvorbis.vorbis_dsp_state;
import org.xiph.libvorbis.vorbis_info;
import org.xiph.libvorbis.vorbisenc;

/**
 * 
 * Limitation: encoding channel count:2; if payload is mono, it is simulated as 2 channels; if more channels provided, only first two channels are encoded.
 * 
 * @since 2016-10-13
 * @author David Zhang (zdl@zdl.hk)
 *
 */
public class VorbisAudioOutputStream extends AbstractAudioOutputStream implements FloatSampleWriter {

    private static final int ENCODING_CHANNELS = 2;
	private float base_quality;

	private final vorbisenc encoder = new vorbisenc();
	private vorbis_info vi = new vorbis_info();
	private vorbis_dsp_state vd = new vorbis_dsp_state();
	private vorbis_block vb;
	private ogg_stream_state os;
	private ogg_page og;
	private ogg_packet op;
	private boolean header_writen = false;

	public VorbisAudioOutputStream(AudioFormat format, long length, TDataOutputStream dataOutputStream, float base_quality) {
		super(dataOutputStream, format, length);
		this.base_quality = base_quality;
		init();
	}

	private void init() {
		if (!encoder.vorbis_encode_init_vbr(vi, ENCODING_CHANNELS, (int) getFormat().getSampleRate(), base_quality)) {
			throw new IllegalArgumentException("Failed to Initialize vorbisenc");
		}
	}

	private int writeHeader() throws IOException {
		vorbis_comment vc = new vorbis_comment();
		vc.vorbis_comment_add_tag("ENCODER", "Java Vorbis Encoder");

		if (!vd.vorbis_analysis_init(vi)) {
			throw new IllegalArgumentException("Failed to Initialize vorbis_dsp_state");
		}

		vb = new vorbis_block(vd);
		os = new ogg_stream_state(new java.util.Random().nextInt());

		ogg_packet header = new ogg_packet();
		ogg_packet header_comm = new ogg_packet();
		ogg_packet header_code = new ogg_packet();

		vd.vorbis_analysis_headerout(vc, header, header_comm, header_code);

		os.ogg_stream_packetin(header); // automatically placed in its own page
		os.ogg_stream_packetin(header_comm);
		os.ogg_stream_packetin(header_code);

		og = new ogg_page();
		op = new ogg_packet();

		if (!os.ogg_stream_flush(og))
			throw new IOException();

		dataOutputStream.write(og.header, 0, og.header_len);
		dataOutputStream.write(og.body, 0, og.body_len);
		header_writen = true;
		return og.header_len + og.body_len;
	}

	@Override
	public void write(FloatSampleBuffer buffer) throws IOException {
		write(buffer, 0, buffer.getSampleCount());
	}

	@Override
	public void write(FloatSampleBuffer fsb, int offset, int sampleCount) throws IOException {
		doWriteBuffer(fsb, offset, sampleCount);
	}

	private int doWriteBuffer(FloatSampleBuffer fsb, int offset, int sampleCount) throws IOException {
		int bytes_writen = 0;
		if (!header_writen) {
			bytes_writen += writeHeader();
		}

		float[][] buffer = vd.vorbis_analysis_buffer(sampleCount*ENCODING_CHANNELS*2);

		if (fsb.getChannelCount() >= ENCODING_CHANNELS) {
			// uninterleave samples
			for (int i = 0; i < ENCODING_CHANNELS; i++) {
				System.arraycopy(fsb.getChannel(i), 0, buffer[i], vd.pcm_current, fsb.getSampleCount());
			}
		} else {
			// Mono Fix
			for (int i = 0; i < ENCODING_CHANNELS; i++) {
				System.arraycopy(fsb.getChannel(0), 0, buffer[i], vd.pcm_current, fsb.getSampleCount());
			}
		}
		// tell the library how much we actually submitted
		vd.vorbis_analysis_wrote(sampleCount);

		// vorbis does some data preanalysis, then divvies up blocks for more involved (potentially parallel) processing. Get a single block for encoding now
		bytes_writen += flush_packet();
		return bytes_writen;
	}

	@Override
	public int write(byte[] abData, int nOffset, int nLength) throws IOException {
		FloatSampleBuffer fsb = new FloatSampleBuffer(abData, nOffset, nLength, getFormat());
		return doWriteBuffer(fsb, 0, fsb.getSampleCount());
	}

	private int flush_packet() throws IOException {
		int bytes_writen = 0;

		while (vb.vorbis_analysis_blockout(vd)) {

			// analysis, assume we want to use bitrate management

			vb.vorbis_analysis(null);
			vb.vorbis_bitrate_addblock();

			while (vd.vorbis_bitrate_flushpacket(op)) {

				// weld the packet into the bitstream
				os.ogg_stream_packetin(op);

				// write out pages (if any)
				while (os.ogg_stream_pageout(og)) {

					dataOutputStream.write(og.header, 0, og.header_len);
					dataOutputStream.write(og.body, 0, og.body_len);

					bytes_writen += og.header_len + og.body_len;

					// this could be set above, but for illustrative purposes, I do it here (to show that vorbis does know where the stream ends)
					if (og.ogg_page_eos() > 0)
						break;
				}
			}
		}
		return bytes_writen;
	}

	@Override
	public void close() throws IOException {
		vd.vorbis_analysis_wrote(0);
		flush_packet();
		dataOutputStream.close();
	}

}